import os

from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

from langchain_core.tools import tool

# 初始化大模型（Agent的"大脑"）
llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    temperature=0  # 0表示输出更确定，适合需要精确决策的场景
)


# 定义工具：计算器（只能做数学运算）
@tool(
    description="用于计算加减乘除运算.参数:expression(必填,数学表达式,示例: 2+4,10-2,3*6,10/2,2+3*7),返回值:计算结果.仅处理数字和+、-、*、/")
def CalculatorTool(expression: str) -> str:
    res = eval(expression)
    return f"计算结果:{expression}={res}"


tools = [CalculatorTool]
from langchain.agents import ZeroShotAgent, AgentExecutor

prompt_template = """
你是一个帮助用户解决问题的智能助手，你拥有以下工具:
{tools}

使用规则:
1:先理解问题:无需工具则直接回答,需要工具则调用工具(根据工具的描述匹配)
2:调用工具必须用格式：```json{{"name":"工具名","parameters":{{"参数名":"参数值"}}}}```（参数需齐全）；
3:拿到工具结果后, 判断是否继续调用工具, 如果需要继续调用工具,则调用工具,否则整理,返回结果;


用户的问题:{input}
"""

prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["tools", "input"],
)

# 创建Agent
agent = ZeroShotAgent.from_llm_and_tools(
    llm=llm,
    tools=tools,
    prompt=prompt,
    verbose=True,
)

# 构建执行器
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=5,
)

# 让Agent解决数学问题
result = agent_executor.invoke({"input": "2*199的结果是啥"})
print("答案：", result)
